﻿/***********************/
/* (c) Alexandru Lungu */
/*  Challenge is Life! */
/* www.challenge-me.ws */
/***********************/

using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using Coowork.WcfCommon;
using System.Xml;
using System.Security.Cryptography.Xml;
using System.Security.Cryptography;
using System.IO;
using System.ServiceModel.Channels;

namespace Coowork.WcfClient
{
    /// <summary>
    /// Provides encryption/decryption on the client.
    /// </summary>
    public class ClientCryptographer : Cryptographer
    {
        //[8]
        public static void Encrypt(XmlDocument xmlDoc, string elementToEncrypt)
        {
            XmlNodeList elementsToEncrypt = xmlDoc.GetElementsByTagName(elementToEncrypt);
            if (elementsToEncrypt.Count == 0)
                return;

            AesCryptoServiceProvider aesServiceProvider = new AesCryptoServiceProvider();
            aesServiceProvider.KeySize = 256;
            aesServiceProvider.GenerateKey();

            XmlNode idNode = xmlDoc.GetElementsByTagName("a:MessageID")[0];

            string id = idNode.InnerText;
            AesKeys.TryAdd(id, aesServiceProvider.Key);

            XmlElement xmlElementToEncrypt = (XmlElement)elementsToEncrypt[0];
            EncryptedXml encryptedXml = new EncryptedXml();

            byte[] encryptedElement = encryptedXml.EncryptData(xmlElementToEncrypt, aesServiceProvider, Content);

            EncryptedData encryptedData = new EncryptedData();
            encryptedData.Type = EncryptedXml.XmlEncElementUrl;
            encryptedData.EncryptionMethod = new EncryptionMethod(EncryptedXml.XmlEncAES256Url);
            EncryptedKey encryptedKey = new EncryptedKey();
            encryptedKey.CipherData = new CipherData(EncryptedXml.EncryptKey(aesServiceProvider.Key, RsaServiceProvider, Content));
            encryptedKey.EncryptionMethod = new EncryptionMethod(EncryptedXml.XmlEncRSA15Url);
            encryptedData.KeyInfo = new KeyInfo();
            encryptedKey.KeyInfo.AddClause(new KeyInfoName(KeyElementName));
            encryptedData.KeyInfo.AddClause(new KeyInfoEncryptedKey(encryptedKey));
            encryptedData.CipherData.CipherValue = encryptedElement;
            encryptedData.Id = id;
            EncryptedXml.ReplaceElement(xmlElementToEncrypt, encryptedData, Content);
        }

        //[8]
        public static void Decrypt(XmlDocument xmlDoc)
        {

            XmlNodeList encryptedElements = xmlDoc.GetElementsByTagName("EncryptedData");
            if (encryptedElements.Count == 0)
                return;

            AesCryptoServiceProvider aesServiceProvider = new AesCryptoServiceProvider();
            XmlElement encryptedElement = (XmlElement)encryptedElements[0];

            EncryptedData encryptedData = new EncryptedData();
            encryptedData.LoadXml(encryptedElement);


            byte[] key;
            AesKeys.TryRemove(encryptedData.Id, out key);
            aesServiceProvider.Key = key;


            EncryptedXml encryptedXml = new EncryptedXml();
            encryptedXml.ReplaceData(encryptedElement, encryptedXml.DecryptData(encryptedData, aesServiceProvider));
        }

        public static ArraySegment<byte> EncryptBuffer(ArraySegment<byte> buffer, BufferManager bufferManager, int messageOffset, string elementToEncrypt = "s:Envelope")
        {
            byte[] bufferedBytes;
            byte[] encryptedBytes;
            XmlDocument xmlDoc = new XmlDocument();

            using (MemoryStream memoryStream = new MemoryStream(buffer.Array, buffer.Offset, buffer.Count))
            {
                xmlDoc.Load(memoryStream);
            }

            ClientCryptographer.Encrypt(xmlDoc, elementToEncrypt);
            encryptedBytes = Encoding.UTF8.GetBytes(xmlDoc.OuterXml);
            bufferedBytes = bufferManager.TakeBuffer(encryptedBytes.Length);
            Array.Copy(encryptedBytes, 0, bufferedBytes, 0, encryptedBytes.Length);
            bufferManager.ReturnBuffer(buffer.Array);

            ArraySegment<byte> byteArray = new ArraySegment<byte>(bufferedBytes, messageOffset, encryptedBytes.Length);



            return byteArray;
        }

        public static ArraySegment<byte> DecryptBuffer(ArraySegment<byte> buffer, BufferManager bufferManager)
        {

            byte[] bufferedBytes;
            ArraySegment<byte> byteArray;
            XmlDocument xmlDoc = new XmlDocument();
            using (MemoryStream memoryStream = new MemoryStream(buffer.Array, buffer.Offset, buffer.Count))
            {
                xmlDoc.Load(memoryStream);
            }

            ClientCryptographer.Decrypt(xmlDoc);
            byte[] decryptedBytes = Encoding.UTF8.GetBytes(xmlDoc.OuterXml);

            bufferedBytes = bufferManager.TakeBuffer(decryptedBytes.Length);
            Array.Copy(decryptedBytes, 0, bufferedBytes, 0, decryptedBytes.Length);
            bufferManager.ReturnBuffer(buffer.Array);
            byteArray = new ArraySegment<byte>(bufferedBytes, 0, decryptedBytes.Length);

            string s = Encoding.UTF8.GetString(byteArray.Array, byteArray.Offset, byteArray.Count);
            return byteArray;

        }

    }

}
